#!/usr/bin/env python3
"""
数据库初始化脚本
用于创建所有定义的表结构
"""

import sys
import os
from contextlib import contextmanager

# 将项目根目录添加到Python路径中
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from sqlalchemy import text
from app.config.database import engine, Base
from app.models.user import User
from app.models.student import Student
from app.models.teacher import Teacher
from app.models.classroom import Classroom
from app.models.course import Course
from app.models.academic_year import AcademicYear
from app.models.class_model import Class
from app.models.class_student import ClassStudent
from app.models.role import Role
from app.models.user_role import UserRole
from app.models.permission import Permission
from app.models.role_permission import RolePermission
from app.models.teacher import Teacher
from app.models.classroom import Classroom
from app.models.course import Course
from app.models.academic_year import AcademicYear
from app.models.class_model import Class
from app.models.class_student import ClassStudent

@contextmanager
def get_connection():
    """获取数据库连接的上下文管理器"""
    connection = engine.connect()
    try:
        yield connection
    finally:
        connection.close()

def check_db_connection():
    """检查数据库连接"""
    try:
        with get_connection() as connection:
            result = connection.execute(text("SELECT 1"))
            print("✓ 数据库连接成功")
            return True
    except Exception as e:
        print(f"✗ 数据库连接失败: {e}")
        return False

def get_existing_tables():
    """获取当前数据库中已存在的表"""
    try:
        # 使用SQLAlchemy的inspect功能获取表信息
        from sqlalchemy import inspect
        inspector = inspect(engine)
        return set(inspector.get_table_names())
    except Exception as e:
        print(f"无法获取现有表信息: {e}")
        return set()

def init_db():
    """初始化数据库表结构"""
    print("开始初始化数据库...")
    
    # 检查数据库连接
    if not check_db_connection():
        print("无法连接到数据库，请检查配置")
        return False
    
    # 获取当前已存在的表
    existing_tables = get_existing_tables()
    print(f"当前已存在的表: {existing_tables}")
    
    # 导入所有模型后创建所有表
    Base.metadata.create_all(bind=engine)
    
    # 获取创建后的所有表
    all_tables = get_existing_tables()
    new_tables = all_tables - existing_tables
    
    print("数据库表结构初始化完成！")
    if new_tables:
        print("新创建的表:")
        for table in new_tables:
            print(f"  + {table}")
    else:
        print("没有创建新表，所有表已存在")
        
    print("\n当前数据库中的所有表:")
    for table in all_tables:
        print(f"  - {table}")
        
    return True

def drop_all_tables():
    """删除所有表（谨慎使用）"""
    print("警告：这将删除所有表及其数据！")
    confirm = input("请输入 'YES' 确认删除所有表: ")
    if confirm == 'YES':
        print("开始删除所有表...")
        Base.metadata.drop_all(bind=engine)
        print("所有表已删除")
        return True
    else:
        print("操作已取消")
        return False

def main():
    """主函数"""
    if len(sys.argv) > 1:
        if sys.argv[1] == 'drop':
            drop_all_tables()
        else:
            print("用法: python init_db.py [drop]")
            print("  drop: 删除所有表")
    else:
        init_db()

if __name__ == "__main__":
    main()